import os
import sys
import logging
from pathlib import Path
from typing import Optional
from contextlib import asynccontextmanager
from fastapi import FastAPI, File, UploadFile, HTTPException, Form
import base64
import io
from PIL import Image
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from rag_engine.rag_chain import RAGEngine
from rag_engine.persona_manager import persona_manager

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Pydantic模型
class QueryRequest(BaseModel):
    question: str
    k: int = 3

class QueryResponse(BaseModel):
    question: str
    answer: str
    sources: list
    status: str = "success"

class UploadResponse(BaseModel):
    filename: str
    status: str
    message: str
    document_id: str

class ImageQueryRequest(BaseModel):
    question: str
    image_base64: str
    k: int = 3

class ImageQueryResponse(BaseModel):
    question: str
    answer: str
    status: str = "success"

class PersonaRequest(BaseModel):
    persona_id: str

class PersonaResponse(BaseModel):
    status: str
    message: str
    persona: dict = None

class PersonaListResponse(BaseModel):
    status: str
    personas: dict
    current_persona: str = None

class PersonaCreateRequest(BaseModel):
    persona_id: str
    name: str
    description: str
    personality: str
    background: str
    speaking_style: str
    prompt_template: str
    fallback_template: str

class SimilarityThresholdRequest(BaseModel):
    threshold: float

class ASRData(BaseModel):
    """ASR数据模型"""
    start: int  # 开始时间戳
    end: int    # 结束时间戳
    text: str   # 识别的文本内容
    confidence: float  # 置信度
    keyframe: Optional[str] = None  # 关键帧信息

class ASRResponse(BaseModel):
    """ASR接收响应模型"""
    status: str
    message: str
    timestamp: str
    data_id: Optional[str] = None

# 创建FastAPI应用
@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    # 启动时初始化
    global rag_engine
    try:
        # 创建配置
        config = {
            'documents_dir': './documents',
            'vector_store_dir': './chroma_db',
            'local_llm_path': './models/Qwen2.5-VL-3B-Instruct',
            'embedding_model_path': './models/bge-large-zh-v1.5',
            'text_splitter_type': 'recursive',
            'chunk_size': 1000,
            'chunk_overlap': 200
        }
        rag_engine = RAGEngine(config=config)
        logger.info("RAG引擎初始化成功")
    except Exception as e:
        logger.error(f"RAG引擎初始化失败: {e}")
    yield
    # 关闭时的清理工作
    logger.info("应用关闭")

app = FastAPI(title="RAG系统", description="文档问答系统", version="1.0.0", lifespan=lifespan)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局变量
rag_engine = None
UPLOAD_DIR = Path("uploads")
UPLOAD_DIR.mkdir(exist_ok=True)



@app.get("/")
async def root():
    """根路径"""
    return {"message": "RAG系统运行正常", "status": "success"}

@app.get("/health")
async def health_check():
    """健康检查"""
    if rag_engine is None:
        return {"status": "unhealthy", "message": "RAG引擎未初始化"}
    return {"status": "healthy", "message": "系统运行正常"}

@app.post("/query", response_model=QueryResponse)
async def query_documents(request: QueryRequest):
    """查询文档"""
    if not rag_engine:
        raise HTTPException(status_code=503, detail="RAG引擎未初始化")
    
    try:
        result = rag_engine.query(
            question=request.question,
            k=request.k
        )
        
        if result['status'] == 'error':
            raise HTTPException(status_code=500, detail=result['error'])
        
        return QueryResponse(**result)
        
    except Exception as e:
        logger.error(f"查询失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# 文档上传功能已禁用 - 系统使用预设的《神曲》文档
# @app.post("/upload", response_model=UploadResponse)
# async def upload_document(file: UploadFile = File(...)):
#     """上传文档 - 已禁用"""
#     raise HTTPException(status_code=403, detail="文档上传功能已禁用，系统使用预设的《神曲》文档")

@app.get("/system/info")
async def get_system_info():
    """获取系统信息"""
    return {
        "system_name": "但丁数字自传系统",
        "persona": "但丁·阿利吉耶里",
        "document": "《神曲》",
        "description": "专门的但丁角色数字自传服务，基于《神曲》文档进行智能问答",
        "features": [
            "但丁角色扮演",
            "《神曲》文档检索",
            "智能问答",
            "ASR数据接收"
        ],
        "status": "ready"
    }

@app.post("/image_query", response_model=ImageQueryResponse)
async def image_query(request: ImageQueryRequest):
    """图片问答接口"""
    try:
        from rag_engine.local_llm import local_llm_manager
        from rag_engine.local_llm import Qwen25VL
        
        # 解码base64图片
        image_data = base64.b64decode(request.image_base64)
        image = Image.open(io.BytesIO(image_data))
        
        # 获取多模态模型
        model = local_llm_manager.get_current_model()
        if not model or not isinstance(model, Qwen25VL):
            raise HTTPException(status_code=503, detail="本地多模态模型不可用")
        
        # 使用专门的图片处理方法
        answer = model.process_image(image, request.question)
        
        return ImageQueryResponse(
            question=request.question,
            answer=answer
        )
        
    except Exception as e:
        logger.error(f"图片查询失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 但丁角色信息API（只读）
@app.get("/dante/info")
async def get_dante_info():
    """获取但丁角色信息"""
    try:
        dante_info = persona_manager.get_persona_info("dante")
        
        return {
            "status": "success",
            "persona": dante_info,
            "system_mode": "dante_digital_biography",
            "document": "《神曲》",
            "capabilities": [
                "基于《神曲》的智能问答",
                "但丁角色扮演",
                "文学作品解读",
                "历史背景介绍"
            ]
        }
    except Exception as e:
        logger.error(f"获取但丁信息失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/dante/threshold", response_model=PersonaResponse)
async def set_similarity_threshold(request: SimilarityThresholdRequest):
    """设置相似度阈值"""
    try:
        persona_manager.set_similarity_threshold(request.threshold)
        
        return PersonaResponse(
            status="success",
            message=f"相似度阈值已设置为: {request.threshold}"
        )
    except Exception as e:
        logger.error(f"设置相似度阈值失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# ASR数据接收API端点
@app.post("/asr/data", response_model=ASRResponse)
async def receive_asr_data(asr_data: ASRData):
    """
    接收ASR服务推送的数据
    
    Args:
        asr_data: ASR识别数据，包含文本和关键帧信息
        
    Returns:
        ASRResponse: 处理结果
    """
    try:
        from datetime import datetime
        
        logger.info(f"📢 收到ASR数据: {asr_data.text[:50]}...")
        
        # 数据验证
        if not asr_data.text or not asr_data.text.strip():
            raise HTTPException(status_code=400, detail="文本内容不能为空")
        
        if asr_data.confidence < 0 or asr_data.confidence > 1:
            raise HTTPException(status_code=400, detail="置信度必须在0-1之间")
        
        # 生成数据ID
        data_id = f"asr_{int(datetime.now().timestamp() * 1000)}"
        
        # 处理ASR数据 - 可以与RAG系统集成
        text_content = asr_data.text.strip()
        
        # 可选：将ASR文本发送给RAG系统进行智能问答
        if rag_engine and len(text_content) > 10:  # 只处理有意义的文本
            try:
                # 使用但丁角色处理ASR文本
                rag_result = rag_engine.query(text_content, k=3, use_persona=True)
                logger.info(f"🎭 RAG处理结果: {rag_result.get('answer', '')[:100]}...")
            except Exception as e:
                logger.warning(f"RAG处理ASR文本失败: {str(e)}")
        
        logger.info(f"✅ ASR数据处理完成，ID: {data_id}")
        
        return ASRResponse(
            status="success",
            message="ASR数据接收成功",
            timestamp=datetime.now().isoformat(),
            data_id=data_id
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"❌ 处理ASR数据失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")


@app.get("/asr/health")
async def asr_health_check():
    """
    ASR服务健康检查
    """
    from datetime import datetime
    return {
        "status": "healthy",
        "service": "ASR数据接收服务",
        "timestamp": datetime.now().isoformat(),
        "rag_engine_status": "available" if rag_engine else "unavailable"
    }


@app.post("/asr/batch", response_model=ASRResponse)
async def receive_asr_batch(asr_data_list: list[ASRData]):
    """
    批量接收ASR数据
    
    Args:
        asr_data_list: ASR数据列表
        
    Returns:
        ASRResponse: 批量处理结果
    """
    try:
        from datetime import datetime
        
        if not asr_data_list:
            raise HTTPException(status_code=400, detail="数据列表不能为空")
        
        if len(asr_data_list) > 100:  # 限制批量大小
            raise HTTPException(status_code=400, detail="批量数据不能超过100条")
        
        processed_count = 0
        failed_count = 0
        
        for asr_data in asr_data_list:
            try:
                text_content = asr_data.text.strip()
                if text_content:
                    logger.info(f"📢 处理ASR文本: {text_content[:30]}...")
                    processed_count += 1
                else:
                    failed_count += 1
            except Exception as e:
                logger.error(f"处理单条ASR数据失败: {str(e)}")
                failed_count += 1
        
        batch_id = f"batch_{int(datetime.now().timestamp() * 1000)}"
        
        return ASRResponse(
            status="success" if failed_count == 0 else "partial_success",
            message=f"批量处理完成: 成功{processed_count}条, 失败{failed_count}条",
            timestamp=datetime.now().isoformat(),
            data_id=batch_id
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"❌ 批量处理ASR数据失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"批量处理失败: {str(e)}")


if __name__ == "__main__":
    import uvicorn
    logger.info("🚀 启动RAG系统服务器...")
    logger.info("📡 服务将在以下地址启动:")
    logger.info("   - 本地访问: http://localhost:8000")
    logger.info("   - 网络访问: http://192.168.31.152:8000")
    logger.info("   - ASR接口: http://192.168.31.152:8000/asr/data")
    uvicorn.run(
        app,
        host="0.0.0.0",  # 监听所有网络接口
        port=8000,
        log_level="info"
    )